from abc import ABCMeta, abstractmethod
from typing import Tuple, Optional, Sequence
from torch import Tensor
import torch

from alphagen.data.expression import Expression
from alphagen.utils.correlation import batch_pearsonr, batch_spearmanr


class AlphaCalculator(metaclass=ABCMeta):
    """Alpha因子计算器的抽象基类
    提供计算单个Alpha因子、多个Alpha因子组合的IC和RankIC的接口
    """
    
    @abstractmethod
    def calc_single_IC_ret(self, expr: Expression) -> float:
        '计算单个Alpha因子与预定义目标之间的信息系数(IC)'

    @abstractmethod
    def calc_single_rIC_ret(self, expr: Expression) -> float:
        '计算单个Alpha因子与预定义目标之间的排序信息系数(Rank IC)'

    def calc_single_all_ret(self, expr: Expression) -> Tuple[float, float]:
        '同时计算单个Alpha因子的IC和RankIC'
        return self.calc_single_IC_ret(expr), self.calc_single_rIC_ret(expr)

    @abstractmethod
    def calc_mutual_IC(self, expr1: Expression, expr2: Expression) -> float:
        '计算两个Alpha因子之间的信息系数'

    @abstractmethod
    def calc_pool_IC_ret(self, exprs: Sequence[Expression], weights: Sequence[float]) -> float:
        '先将多个Alpha因子线性组合，然后计算组合Alpha与预定义目标之间的IC'

    @abstractmethod
    def calc_pool_rIC_ret(self, exprs: Sequence[Expression], weights: Sequence[float]) -> float:
        '先将多个Alpha因子线性组合，然后计算组合Alpha与预定义目标之间的RankIC'

    @abstractmethod
    def calc_pool_all_ret(self, exprs: Sequence[Expression], weights: Sequence[float]) -> Tuple[float, float]:
        '先将多个Alpha因子线性组合，然后同时计算组合Alpha的IC和RankIC'


class TensorAlphaCalculator(AlphaCalculator):
    """基于张量计算的Alpha因子计算器实现类"""
    
    def __init__(self, target: Optional[Tensor]) -> None:
        '初始化计算器，设置预定义的目标收益率'
        self._target = target

    @property
    @abstractmethod
    def n_days(self) -> int: 
        '获取数据的交易日天数'

    @property
    def target(self) -> Tensor:
        '获取目标收益率张量'
        if self._target is None:
            raise ValueError("在计算非互信息系数前必须设置目标收益率")
        return self._target

    @abstractmethod
    def evaluate_alpha(self, expr: Expression) -> Tensor:
        '将Alpha表达式计算为形状为(days, stocks)的张量'

    def make_ensemble_alpha(self, exprs: Sequence[Expression], weights: Sequence[float]) -> Tensor:
        '将多个Alpha因子进行加权组合'
        n = len(exprs)
        factors = [self.evaluate_alpha(exprs[i]) * weights[i] for i in range(n)]
        return torch.sum(torch.stack(factors, dim=0), dim=0)

    def _calc_IC(self, value1: Tensor, value2: Tensor) -> float:
        '计算两个张量之间的Pearson相关系数均值'
        return batch_pearsonr(value1, value2).mean().item()
    
    def _calc_rIC(self, value1: Tensor, value2: Tensor) -> float:
        '计算两个张量之间的Spearman相关系数均值'
        return batch_spearmanr(value1, value2).mean().item()
    
    def _IR_from_batch(self, batch: Tensor) -> float:
        '计算信息比率(均值/标准差)'
        mean, std = batch.mean(), batch.std()
        return (mean / std).item()
    
    def _calc_ICIR(self, value1: Tensor, value2: Tensor) -> float:
        '计算IC的信息比率'
        return self._IR_from_batch(batch_pearsonr(value1, value2))
    
    def _calc_rICIR(self, value1: Tensor, value2: Tensor) -> float:
        '计算RankIC的信息比率'
        return self._IR_from_batch(batch_spearmanr(value1, value2))

    def calc_single_IC_ret(self, expr: Expression) -> float:
        return self._calc_IC(self.evaluate_alpha(expr), self.target)
    
    def calc_single_IC_ret_daily(self, expr: Expression) -> Tensor:
        '计算单个Alpha因子的每日IC'
        return batch_pearsonr(self.evaluate_alpha(expr), self.target)

    def calc_single_rIC_ret(self, expr: Expression) -> float:
        return self._calc_rIC(self.evaluate_alpha(expr), self.target)
    
    def calc_single_all_ret(self, expr: Expression) -> Tuple[float, float]:
        value = self.evaluate_alpha(expr)
        target = self.target
        return self._calc_IC(value, target), self._calc_rIC(value, target)

    def calc_mutual_IC(self, expr1: Expression, expr2: Expression) -> float:
        return self._calc_IC(self.evaluate_alpha(expr1), self.evaluate_alpha(expr2))

    def calc_mutual_IC_daily(self, expr1: Expression, expr2: Expression) -> Tensor:
        '计算两个Alpha因子之间的每日IC'
        return batch_pearsonr(self.evaluate_alpha(expr1), self.evaluate_alpha(expr2))

    def calc_pool_IC_ret(self, exprs: Sequence[Expression], weights: Sequence[float]) -> float:
        with torch.no_grad():
            value = self.make_ensemble_alpha(exprs, weights)
            return self._calc_IC(value, self.target)

    def calc_pool_rIC_ret(self, exprs: Sequence[Expression], weights: Sequence[float]) -> float:
        with torch.no_grad():
            value = self.make_ensemble_alpha(exprs, weights)
            return self._calc_rIC(value, self.target)

    def calc_pool_all_ret(self, exprs: Sequence[Expression], weights: Sequence[float]) -> Tuple[float, float]:
        with torch.no_grad():
            value = self.make_ensemble_alpha(exprs, weights)
            target = self.target
            return self._calc_IC(value, target), self._calc_rIC(value, target)
        
    def calc_pool_all_ret_with_ir(self, exprs: Sequence[Expression], weights: Sequence[float]) -> Tuple[float, float, float, float]:
        """计算组合Alpha的完整指标
        Returns:
            Tuple: (IC均值, IC信息比率, RankIC均值, RankIC信息比率)
        """
        with torch.no_grad():
            value = self.make_ensemble_alpha(exprs, weights)
            target = self.target
            ics = batch_pearsonr(value, target)
            rics = batch_spearmanr(value, target)
            ic_mean, ic_std = ics.mean().item(), ics.std().item()
            ric_mean, ric_std = rics.mean().item(), rics.std().item()
            return ic_mean, ic_mean / ic_std, ric_mean, ric_mean / ric_std
